package net.minecraft.entity.passive;

import net.minecraft.registry.entry.RegistryEntry;
import net.minecraft.registry.tag.BiomeTags;
import net.minecraft.util.DyeColor;
import net.minecraft.util.collection.Pool;
import net.minecraft.util.math.random.Random;
import net.minecraft.world.biome.Biome;

public class SheepColors {
	private static final SheepColors.SpawnConfig TEMPERATE = new SheepColors.SpawnConfig(
		createCombinedSelector(
			poolBuilder()
				.add(createSingleSelector(DyeColor.BLACK), 5)
				.add(createSingleSelector(DyeColor.GRAY), 5)
				.add(createSingleSelector(DyeColor.LIGHT_GRAY), 5)
				.add(createSingleSelector(DyeColor.BROWN), 3)
				.add(createDefaultSelector(DyeColor.WHITE), 82)
				.build()
		)
	);
	private static final SheepColors.SpawnConfig WARM = new SheepColors.SpawnConfig(
		createCombinedSelector(
			poolBuilder()
				.add(createSingleSelector(DyeColor.GRAY), 5)
				.add(createSingleSelector(DyeColor.LIGHT_GRAY), 5)
				.add(createSingleSelector(DyeColor.WHITE), 5)
				.add(createSingleSelector(DyeColor.BLACK), 3)
				.add(createDefaultSelector(DyeColor.BROWN), 82)
				.build()
		)
	);
	private static final SheepColors.SpawnConfig COLD = new SheepColors.SpawnConfig(
		createCombinedSelector(
			poolBuilder()
				.add(createSingleSelector(DyeColor.LIGHT_GRAY), 5)
				.add(createSingleSelector(DyeColor.GRAY), 5)
				.add(createSingleSelector(DyeColor.WHITE), 5)
				.add(createSingleSelector(DyeColor.BROWN), 3)
				.add(createDefaultSelector(DyeColor.BLACK), 82)
				.build()
		)
	);

	private static SheepColors.ColorSelector createDefaultSelector(DyeColor color) {
		return createCombinedSelector(poolBuilder().add(createSingleSelector(color), 499).add(createSingleSelector(DyeColor.PINK), 1).build());
	}

	public static DyeColor select(RegistryEntry<Biome> biome, Random random) {
		SheepColors.SpawnConfig spawnConfig = getSpawnConfig(biome);
		return spawnConfig.colors().get(random);
	}

	private static SheepColors.SpawnConfig getSpawnConfig(RegistryEntry<Biome> biome) {
		if (biome.isIn(BiomeTags.SPAWNS_WARM_VARIANT_FARM_ANIMALS)) {
			return WARM;
		} else {
			return biome.isIn(BiomeTags.SPAWNS_COLD_VARIANT_FARM_ANIMALS) ? COLD : TEMPERATE;
		}
	}

	private static SheepColors.ColorSelector createCombinedSelector(Pool<SheepColors.ColorSelector> pool) {
		if (pool.isEmpty()) {
			throw new IllegalArgumentException("List must be non-empty");
		} else {
			return random -> pool.get(random).get(random);
		}
	}

	private static SheepColors.ColorSelector createSingleSelector(DyeColor color) {
		return random -> color;
	}

	private static Pool.Builder<SheepColors.ColorSelector> poolBuilder() {
		return Pool.builder();
	}

	@FunctionalInterface
	interface ColorSelector {
		DyeColor get(Random random);
	}

	record SpawnConfig(SheepColors.ColorSelector colors) {
	}
}
